import torch
import torchvision.transforms

from framework import yolov3
from backbone.darknet53 import YoloBody
from utils.dataset import BananaDataset
from utils.object_detection import bbox_position_cast
from utils.train import init_weights

import torch.utils.data

if __name__ == '__main__':
    device = torch.device('cuda')
    max_anchors = torch.tensor([[116, 90],
                                [156, 198],
                                [373, 326]])
    medium_anchors = torch.tensor([[30, 61],
                                   [62, 45],
                                   [59, 119],
                                   ])
    min_anchors = torch.tensor([[10, 13],
                                [16, 30],
                                [33, 23], ])


    yolo_loss = yolov3.YOLOLoss([max_anchors, medium_anchors, min_anchors], c=1, strides=[13, 26, 52],
                                image_size=torch.Size([416, 416]), device=device)

    model = YoloBody([3, 3, 3], 1)

    epoch = 30


    def change_y_pos(y):
        y[..., 1:5] = bbox_position_cast(y[..., 1:5], 1)
        return y


    target_transform = torchvision.transforms.Compose(
        [change_y_pos]
    )

    transform = torchvision.transforms.Resize((416, 416))

    train_dataset = BananaDataset('./out/data/datasets/banana-detection/bananas_train', 'images'
                                  , 'label.csv', transform=transform, target_transform=target_transform)

    model.apply(init_weights)
    model.to(device)

    optimizer = torch.optim.SGD(model.parameters(), 1e-4, weight_decay=5e-4)

    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.94)


    def collate_fn(batch):
        images = []
        labels = []
        for image, label in batch:
            images.append(image)
            labels.append(label)

        return torch.stack(images), labels


    for e in range(epoch):
        model.train()

        for x, y in torch.utils.data.DataLoader(train_dataset, batch_size=8, shuffle=True, collate_fn=collate_fn):
            x = x.to(device)
            for i in range(len(y)):
                y[i] = y[i].to(device)
            optimizer.zero_grad()
            y_p = model(x)
            l = yolo_loss(y_p, y)
            print(l.item())
            l.backward()
            optimizer.step()
